from flask import Flask, request, jsonify
import random
import string
import requests
import sys
import os

app = Flask(__name__)

# 生成随机token的函数
def generate_random_token(length=32):
    charset = string.ascii_letters + string.digits
    return ''.join(random.choice(charset) for _ in range(length))

TOKEN = generate_random_token()
print(f"Generated Token: {TOKEN}")

@app.route('/api/status', methods=['GET'])
def get_version():
    """返回服务器运行状态"""
    auth_header = request.headers.get('Authorization')
    if not auth_header or auth_header != TOKEN:
        return jsonify({"error": "Unauthorized"}), 401

    try:
        # 检查 Ollama 是否运行
        response = requests.get("http://localhost:11434")
        
        # 检查状态码
        if response.status_code == 200:
            return jsonify({
                "status": "Ollama is running"
            }), 200
        elif response.status_code == 404:
            return jsonify({
                "status": "error",
                "error": "Ollama service not found"
            }), 500
        elif response.status_code == 401:
            return jsonify({
                "status": "error",
                "error": "Unauthorized access to Ollama"
            }), 500
        elif response.status_code == 500:
            return jsonify({
                "status": "error",
                "error": "Ollama service error"
            }), 500
        else:
            return jsonify({
                "status": "error",
                "error": f"Unexpected status code: {response.status_code}"
            }), 500
            
    except requests.exceptions.ConnectionError:
        return jsonify({
            "status": "error",
            "error": "Could not connect to Ollama service"
        }), 500
    except requests.exceptions.Timeout:
        return jsonify({
            "status": "error",
            "error": "Connection to Ollama timed out"
        }), 500
    except Exception as e:
        return jsonify({
            "status": "error",
            "error": f"Unexpected error: {str(e)}"
        }), 500

@app.route('/api/generate', methods=['POST'])
def generate_text():
    auth_header = request.headers.get('Authorization')
    if not auth_header or auth_header != TOKEN:
        return jsonify({"error": "Unauthorized"}), 401

    data = request.json
    model = data.get("model")
    prompt = data.get("prompt")
    stream = data.get("stream", False)
    options = data.get("options", {})

    ollama_url = "http://localhost:11434/api/generate"
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": stream,
        "options": options
    }

    try:
        response = requests.post(ollama_url, json=payload)
        response.raise_for_status()
        return response.json(), 200
    except requests.exceptions.RequestException as e:
        return jsonify({"error": f"Failed to call Ollama API: {str(e)}"}), 500

@app.route('/api/chat', methods=['POST'])
def chat():
    auth_header = request.headers.get('Authorization')
    if not auth_header or auth_header != TOKEN:
        return jsonify({"error": "Unauthorized"}), 401

    data = request.json
    model = data.get("model")
    messages = data.get("messages")
    stream = data.get("stream", False)
    options = data.get("options", {})

    ollama_url = "http://localhost:11434/api/chat"
    payload = {
        "model": model,
        "messages": messages,
        "stream": stream,
        "options": options
    }

    try:
        response = requests.post(ollama_url, json=payload)
        response.raise_for_status()
        return response.json(), 200
    except requests.exceptions.RequestException as e:
        return jsonify({"error": f"Failed to call Ollama API: {str(e)}"}), 500

@app.route('/api/embeddings', methods=['POST'])
def embeddings():
    auth_header = request.headers.get('Authorization')
    if not auth_header or auth_header != TOKEN:
        return jsonify({"error": "Unauthorized"}), 401

    data = request.json
    model_name = data.get("model", {}).get("name")
    input_data = data.get("input", [])

    ollama_url = "http://localhost:11434/api/embeddings"
    payload = {
        "model": {
            "name": model_name
        },
        "input": input_data
    }

    try:
        response = requests.post(ollama_url, json=payload)
        response.raise_for_status()
        return response.json(), 200
    except requests.exceptions.RequestException as e:
        return jsonify({"error": f"Failed to call Ollama API: {str(e)}"}), 500

@app.route('/api/chats/<chat_id>/delete', methods=['DELETE'])
def delete_chat(chat_id):
    """删除指定的聊天历史"""
    auth_header = request.headers.get('Authorization')
    if not auth_header or auth_header != TOKEN:
        return jsonify({"error": "Unauthorized"}), 401

    try:
        # 确保 chat_id 是有效的文件名
        if not chat_id or not chat_id.isalnum():
            return jsonify({
                "status": "error",
                "message": "Invalid chat ID"
            }), 400

        history_file = os.path.join(CHAT_HISTORY_DIR, f"{chat_id}.json")
        
        # 检查文件是否存在
        if not os.path.exists(history_file):
            return jsonify({
                "status": "error",
                "message": "Chat not found"
            }), 404

        # 尝试删除文件
        try:
            os.remove(history_file)
            return jsonify({
                "status": "success",
                "message": "Chat deleted successfully"
            })
        except OSError as e:
            return jsonify({
                "status": "error",
                "message": f"Failed to delete chat file: {str(e)}"
            }), 500

    except Exception as e:
        return jsonify({
            "status": "error",
            "message": f"Internal server error: {str(e)}"
        }), 500

if __name__ == '__main__':
    # 默认端口为4433，可以通过命令行参数指定其他端口
    port = 4433
    if len(sys.argv) > 1:
        port = int(sys.argv[1])

    print(f"Server running on port {port}. Press Ctrl+C to stop.")
    app.run(host='0.0.0.0', port=port)